{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47e69cae",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install anndata\n",
    "!pip install scanpy\n",
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "06af81a7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/felixfischer/anaconda3/envs/merlin-torch/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from os.path import join\n",
    "\n",
    "import anndata\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "import torch\n",
    "\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f9ede2f4-6560-4a94-9ee3-2fbceeea5567",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BASE_PATH = '/home/felixfischer/model_inference'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec909451-f531-4eb2-b43b-33b69d8d2c76",
   "metadata": {},
   "source": [
    "# Load example data set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2170b572-bcc7-4fac-90eb-3d40e88ffd4d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 5000 × 24185\n",
       "    obs: 'donor_id', 'nCount_RNA', 'nFeature_RNA', 'percent.mt', 'seurat_clusters', 'celltype', 'organism_ontology_term_id', 'assay_ontology_term_id', 'self_reported_ethnicity_ontology_term_id', 'tissue_ontology_term_id', 'is_primary_data', 'sex_ontology_term_id', 'development_stage_ontology_term_id', 'disease_ontology_term_id', 'cell_type_ontology_term_id', 'suspension_type', 'cell_type', 'assay', 'disease', 'organism', 'sex', 'tissue', 'self_reported_ethnicity', 'development_stage'\n",
       "    var: 'name', 'feature_is_filtered', 'feature_name', 'feature_reference', 'feature_biotype'\n",
       "    uns: 'cell_type_ontology_term_id_colors', 'schema_version', 'seurat_clusters_colors', 'title'\n",
       "    obsm: 'X_pca', 'X_umap'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata = anndata.read_h5ad(join(BASE_PATH, 'local.h5ad'))\n",
    "# subsample to 5000 cells to make inference run faster\n",
    "sc.pp.subsample(adata, n_obs=5000)\n",
    "adata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1143d50f-ec2e-4117-8125-c58c0f3318f7",
   "metadata": {},
   "source": [
    "# Inference steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a16c3721-b1c1-476f-897a-f84a7c0e6a2c",
   "metadata": {},
   "source": [
    "## 1. Preprocess data for model inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eaff8c2-e17d-4f41-9b7d-7bfd825ccb21",
   "metadata": {},
   "source": [
    "Data preparation involves the following steps:\n",
    "\n",
    "1. Streamline feature space (select genes + arrange them in the same order as model was fitted on). \n",
    "    * We should agree on a specific ensembl release\n",
    "        * Currently, I use version 104 - that's the version cellxgene uses\n",
    "        * If we have the same ensembl release we can just do string matching and rearrange the matrix\n",
    "        * If this in not done, the code below does not give the right results\n",
    "2. Wrap data set into pytorch data loader -> use this to feed data into model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c7d0774-b532-47b4-a6f7-0ceef9aa11c4",
   "metadata": {},
   "source": [
    "#### 1. Streamline feature space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c7b77042-f56a-4592-b261-159fef6bd59f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>feature_is_filtered</th>\n",
       "      <th>feature_name</th>\n",
       "      <th>feature_reference</th>\n",
       "      <th>feature_biotype</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>gene_ids</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ENSG00000238009</th>\n",
       "      <td>RP11-34P13.7</td>\n",
       "      <td>False</td>\n",
       "      <td>RP11-34P13.7</td>\n",
       "      <td>NCBITaxon:9606</td>\n",
       "      <td>gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ENSG00000279457</th>\n",
       "      <td>FO538757.2</td>\n",
       "      <td>False</td>\n",
       "      <td>WASH9P</td>\n",
       "      <td>NCBITaxon:9606</td>\n",
       "      <td>gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ENSG00000228463</th>\n",
       "      <td>AP006222.2</td>\n",
       "      <td>False</td>\n",
       "      <td>AP006222.1</td>\n",
       "      <td>NCBITaxon:9606</td>\n",
       "      <td>gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ENSG00000237094</th>\n",
       "      <td>RP4-669L17.10</td>\n",
       "      <td>False</td>\n",
       "      <td>RP4-669L17.4</td>\n",
       "      <td>NCBITaxon:9606</td>\n",
       "      <td>gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ENSG00000230021</th>\n",
       "      <td>RP5-857K21.4</td>\n",
       "      <td>False</td>\n",
       "      <td>RP11-206L10.17</td>\n",
       "      <td>NCBITaxon:9606</td>\n",
       "      <td>gene</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                          name  feature_is_filtered    feature_name  \\\n",
       "gene_ids                                                              \n",
       "ENSG00000238009   RP11-34P13.7                False    RP11-34P13.7   \n",
       "ENSG00000279457     FO538757.2                False          WASH9P   \n",
       "ENSG00000228463     AP006222.2                False      AP006222.1   \n",
       "ENSG00000237094  RP4-669L17.10                False    RP4-669L17.4   \n",
       "ENSG00000230021   RP5-857K21.4                False  RP11-206L10.17   \n",
       "\n",
       "                feature_reference feature_biotype  \n",
       "gene_ids                                           \n",
       "ENSG00000238009    NCBITaxon:9606            gene  \n",
       "ENSG00000279457    NCBITaxon:9606            gene  \n",
       "ENSG00000228463    NCBITaxon:9606            gene  \n",
       "ENSG00000237094    NCBITaxon:9606            gene  \n",
       "ENSG00000230021    NCBITaxon:9606            gene  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.var.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "215010ff-da4e-4b1a-aa0b-6097aa41fb9e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>feature_id</th>\n",
       "      <th>feature_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ENSG00000186092</td>\n",
       "      <td>OR4F5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ENSG00000284733</td>\n",
       "      <td>OR4F29</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ENSG00000284662</td>\n",
       "      <td>OR4F16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ENSG00000187634</td>\n",
       "      <td>SAMD11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ENSG00000188976</td>\n",
       "      <td>NOC2L</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        feature_id feature_name\n",
       "0  ENSG00000186092        OR4F5\n",
       "1  ENSG00000284733       OR4F29\n",
       "2  ENSG00000284662       OR4F16\n",
       "3  ENSG00000187634       SAMD11\n",
       "4  ENSG00000188976        NOC2L"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load the gene order on which model was trained on\n",
    "genes_from_model = pd.read_parquet(join(BASE_PATH, '/home/felixfischer/model_inference/checkpoint/var.parquet'))\n",
    "genes_from_model.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "929ae0c7",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5000, 19331)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from cellnet.utils.data_loading import streamline_count_matrix\n",
    "\n",
    "\n",
    "x_streamlined = streamline_count_matrix(adata.raw.X, adata.var.name, genes_from_model.feature_name)\n",
    "x_streamlined.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dfc8d40-61e8-4ca7-a112-744ca0b91f5a",
   "metadata": {},
   "source": [
    "#### 2. Wrap into pytroch data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ff6e766-014f-49f7-9f69-2939c1d60f24",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from cellnet.utils.data_loading import dataloader_factory\n",
    "\n",
    "# Wrap dataset into pytorch data loader to use for batched inference\n",
    "loader = dataloader_factory(x_streamlined, batch_size=2048)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a44f88a-7515-4706-8520-54e7e477048d",
   "metadata": {},
   "source": [
    "## 2. Load weights from checkpoint and intialize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d51502d4-699d-4f50-973b-2b5728668acd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections import OrderedDict\n",
    "import yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2f54a308-61f7-477f-887d-84d8683bdc5b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load checkpoint\n",
    "ckpt = torch.load(join(BASE_PATH, 'checkpoint/model.ckpt'), map_location=torch.device('cpu'))\n",
    "\n",
    "# extract state_dict of tabnet model from checkpoint\n",
    "# I can do this as well and just send you the updated checkpoint file - I think this would be the best solution\n",
    "# I just put this here for completeness\n",
    "tabnet_weights = OrderedDict()\n",
    "for name, weight in ckpt['state_dict'].items():\n",
    "    if 'classifier.' in name:\n",
    "        tabnet_weights[name.replace('classifier.', '')] = weight\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2b2934ac-4864-4fde-8c64-48d1d38d4ecb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from cellnet.tabnet.tab_network import TabNet\n",
    "\n",
    "\n",
    "# load in hparams file of model to get model architecture\n",
    "with open(join(BASE_PATH, 'checkpoint/hparams.yaml')) as f:\n",
    "    model_params = yaml.full_load(f.read())\n",
    "\n",
    "\n",
    "# initialzie model with hparams from hparams.yaml file\n",
    "tabnet = TabNet(\n",
    "    input_dim=model_params['gene_dim'],\n",
    "    output_dim=model_params['type_dim'],\n",
    "    n_d=model_params['n_d'],\n",
    "    n_a=model_params['n_a'],\n",
    "    n_steps=model_params['n_steps'],\n",
    "    gamma=model_params['gamma'],\n",
    "    n_independent=model_params['n_independent'],\n",
    "    n_shared=model_params['n_shared'],\n",
    "    epsilon=model_params['epsilon'],\n",
    "    virtual_batch_size=model_params['virtual_batch_size'],\n",
    "    momentum=model_params['momentum'],\n",
    "    mask_type=model_params['mask_type'],\n",
    ")\n",
    "\n",
    "# load trained weights\n",
    "tabnet.load_state_dict(tabnet_weights)\n",
    "# set model to inference mode\n",
    "tabnet.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c77f515b-4841-4460-a205-a061309d011d",
   "metadata": {},
   "source": [
    "## 3. Run model inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "05b2ac0e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sf_log1p_norm(x):\n",
    "    \"\"\"Normalize each cell to have 10000 counts and apply log(x+1) transform.\"\"\"\n",
    "\n",
    "    counts = torch.sum(x, dim=1, keepdim=True)\n",
    "    # avoid zero division error\n",
    "    counts += counts == 0.\n",
    "    scaling_factor = 10000. / counts\n",
    "\n",
    "    return torch.log1p(scaling_factor * x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "06c50c15-d2ab-4347-9643-0f80a1a27f21",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:47<00:00, 15.70s/it]\n"
     ]
    }
   ],
   "source": [
    "preds = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch in tqdm(loader):\n",
    "        # normalize data\n",
    "        x_input = sf_log1p_norm(batch[0]['X'])\n",
    "        logits, _ = tabnet(x_input)\n",
    "        preds.append(torch.argmax(logits, dim=1).numpy())\n",
    "\n",
    "\n",
    "preds = np.hstack(preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "da32c5c9-52c6-4290-8942-c5c2a6cc5135",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([118, 127,  22, ..., 118, 107, 118])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "72917ef9-5973-433f-8b32-3f1c7c02fc9b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# model outputs just integers -> each int corresponds to a specific cell type\n",
    "# revert this mapping now\n",
    "cell_type_mapping = pd.read_parquet(join(BASE_PATH, 'checkpoint/cell_type.parquet'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d4629760-fb02-4ce1-a08f-bf191e8cdb3f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['monocyte', 'natural killer cell', 'T cell', ..., 'monocyte',\n",
       "       'mast cell', 'monocyte'], dtype=object)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds = cell_type_mapping.loc[preds]['label'].to_numpy()\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "be65ee0c-5e26-417d-84e2-b428910f22ac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5000,)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preds.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b728bfce-177f-4227-a349-f56924dadf3c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
